package com.demo.minibase;

import java.io.Closeable;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.NavigableMap;
import java.util.SortedMap;
import java.util.concurrent.ConcurrentSkipListMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.ReentrantReadWriteLock;

import org.apache.log4j.Logger;

import com.demo.minibase.DiskStore.MultiIter;
import com.demo.minibase.MStore.SeekIter;
import com.demo.minibase.MiniBase.Flusher;

public class MemStore implements Closeable {

	private static final Logger LOG = Logger.getLogger(MemStore.class);

	/**当前MemStore内存占用字节数，用于判断是否达到Flush阈值。它是一个AtomicLong变量，所以在写入高并发的情景下，该变量可能成为一个性能瓶颈，
	 * <br/>因为所有的并发写入都需要串行执行这个CAS操作。但就目前来说，对性能并不会产生太大的影响
	 */
	private final AtomicLong dataSize = new AtomicLong();

	private volatile ConcurrentSkipListMap<KeyValue, KeyValue> kvMap;
	private volatile ConcurrentSkipListMap<KeyValue, KeyValue> snapshot;

	private final ReentrantReadWriteLock updateLock = new ReentrantReadWriteLock();
	private final AtomicBoolean isSnapshotFlushing = new AtomicBoolean(false);
	private ExecutorService pool;

	private Config conf;
	private Flusher flusher;

	public MemStore(Config conf, Flusher flusher, ExecutorService pool) {
		this.conf = conf;
		this.flusher = flusher;
		this.pool = pool;

		dataSize.set(0);
		this.kvMap = new ConcurrentSkipListMap<>();
		this.snapshot = null;
	}

	public void add(KeyValue kv) throws IOException {
		// 需要考虑的是，数据量达到一个阈值之后，需要开始flush
		flushIfNeeded(true);
		// 在flush的时候，先将MutableMemstore切换成ImmutableMemstore，切换过程中必须保证没有写入操作。所以，我们用一个读写锁updateLock来控制写入操作和Flush操作的互斥。
		// 写入时，先拿到updateLock的读锁，写完之后释放，而在切换MemStore时拿到updateLock的写锁，切换完之后释放
		updateLock.readLock().lock();
		try {
			KeyValue prevKeyValue;
			if ((prevKeyValue = kvMap.put(kv, kv)) == null) {
				// 在成功地把KV放入到ConcurrentSkipListMap之后，需要更新MemStore的dataSize
				// 之前KV不存在，则直接累计kv.getSerializeSize即可
				dataSize.addAndGet(kv.getSerializeSize());
			} else {
				// 之前KV存在，则将二者差值累积到dataSize上
				dataSize.addAndGet(kv.getSerializeSize() - prevKeyValue.getSerializeSize());
			}
		} finally {
			updateLock.readLock().unlock();
		}
		flushIfNeeded(false);
	}

	private void flushIfNeeded(boolean shouldBlocking) throws IOException {
		if (getDataSize() > conf.getMaxMemstoreSize()) {
			if (isSnapshotFlushing.get() && shouldBlocking) {
				// 如果shouldBlocking为true，在发现MemStore满的时候，将抛出IOException，告诉用户MemStore已经满了，此刻重试即可
				throw new IOException("Memstore is full, currentDataSize=" + dataSize.get() + "B, maxMemstoreSize="
						+ conf.getMaxMemstoreSize() + "B, please wait until the flushing is finished.");
			} else if (isSnapshotFlushing.compareAndSet(false, true)) {
				pool.submit(new FlusherTask());
			}
		}
	}

	public long getDataSize() {
		return dataSize.get();
	}

	public boolean isFlushing() {
		return this.isSnapshotFlushing.get();
	}

	@Override
	public void close() throws IOException {
	}

	private class FlusherTask implements Runnable {
		@Override
		public void run() {
			// Step.1 memstore snpashot
			updateLock.writeLock().lock();
			try {
				// snapshot指向kvMap
				snapshot = kvMap;
				// 清空kvMap，所以在scan时，需要把kvMap和snapshot同时扫描
				kvMap = new ConcurrentSkipListMap<>();
				dataSize.set(0);
			} finally {
				updateLock.writeLock().unlock();
			}

			// Step.2 Flush the memstore to disk file.
			boolean success = false;
			for (int i = 0; i < conf.getFlushMaxRetries(); i++) {
				try {
					flusher.flush(new IteratorWrapper(snapshot));
					success = true;
				} catch (IOException e) {
					LOG.error("Failed to flush memstore, retries=" + i + ", maxFlushRetries=" + conf.getFlushMaxRetries(), e);
					if (i >= conf.getFlushMaxRetries()) {
						break;
					}
				}
			}

			// Step.3 clear the snapshot.
			if (success) {
				// TODO MemStoreIter may get a NPE because we set null here ? should synchronize
				// ?
				snapshot = null;
				isSnapshotFlushing.compareAndSet(true, false);
			}
		}
	}

	public SeekIter<KeyValue> createIterator() throws IOException {
		return new MemStoreIter(kvMap, snapshot);
	}

	public static class IteratorWrapper implements SeekIter<KeyValue> {

		private SortedMap<KeyValue, KeyValue> sortedMap;
		private Iterator<KeyValue> it;

		public IteratorWrapper(SortedMap<KeyValue, KeyValue> sortedMap) {
			this.sortedMap = sortedMap;
			this.it = sortedMap.values().iterator();
		}

		@Override
		public boolean hasNext() throws IOException {
			return it != null && it.hasNext();
		}

		@Override
		public KeyValue next() throws IOException {
			return it.next();
		}

		@Override
		public void seekTo(KeyValue kv) throws IOException {
			// 改变原来的迭代器，截取了从kv开始往后的部分
			it = sortedMap.tailMap(kv).values().iterator();
		}
	}

	private class MemStoreIter implements SeekIter<KeyValue> {

		private MultiIter it;

		public MemStoreIter(NavigableMap<KeyValue, KeyValue> kvSet, NavigableMap<KeyValue, KeyValue> snapshot) throws IOException {
			List<IteratorWrapper> inputs = new ArrayList<>();
			if (kvSet != null && kvSet.size() > 0) {
				inputs.add(new IteratorWrapper(kvMap));
			}
			if (snapshot != null && snapshot.size() > 0) {
				inputs.add(new IteratorWrapper(snapshot));
			}
			it = new MultiIter(inputs.toArray(new IteratorWrapper[0]));
		}

		@Override
		public boolean hasNext() throws IOException {
			return it.hasNext();
		}

		@Override
		public KeyValue next() throws IOException {
			return it.next();
		}

		@Override
		public void seekTo(KeyValue kv) throws IOException {
			it.seekTo(kv);
		}
	}
}
